#ifndef PROVIDENT_BINOMIAL_TREE_HPP_
#define PROVIDENT_BINOMIAL_TREE_HPP_

#include <cmath>
#include <cstddef>

#include "VanillaOption.hpp"

namespace provident
{
    class CoxRossRubinstein
    {
    public:
        static void factors(double& upFactor,
                            double& downFactor,
                            double& upProbability,
                            double const volatility,
                            double const riskFreeRate,
                            double const timeStep)
        {
            upFactor = exp(volatility * sqrt(timeStep));
            downFactor = 1.0 / upFactor;
            upProbability = (exp(riskFreeRate*timeStep) - downFactor) / (upFactor - downFactor);
        }
    };

    class EqualProbabilities
    {
    public:
        static void factors(double& upFactor,
                            double& downFactor,
                            double& upProbability,
                            double const volatility,
                            double const riskFreeRate,
                            double const timeStep)
        {
            double const exp2VolSqrtTime(exp(2*volatility*sqrt(timeStep)));
            downFactor = (2*exp(riskFreeRate*timeStep)) / (exp2VolSqrtTime+1);
            upFactor = exp2VolSqrtTime * downFactor;
            upProbability = 0.5;
        }
    };
    
    class JarrowRudd
    {
    public:
        static void factors(double& upFactor,
                            double& downFactor,
                            double& upProbability,
                            double const volatility,
                            double const riskFreeRate,
                            double const timeStep)
        {
            double const a(timeStep*(riskFreeRate - 0.5*volatility*volatility));
            double const b(volatility * sqrt(timeStep));
            upFactor = exp(a+b);
            downFactor = exp(a-b);
            upProbability = 0.5;
        }
    };
    
    class Tian
    {
    public:
        static void factors(double& upFactor,
                            double& downFactor,
                            double& upProbability,
                            double const volatility,
                            double const riskFreeRate,
                            double const timeStep)
        {
            double const a(0.5*exp(timeStep*(riskFreeRate+volatility*volatility)));
            double const b(exp(volatility*volatility*timeStep)+1);
            double const c(sqrt(exp(pow(volatility,4)*timeStep*timeStep)+2*exp(volatility*volatility*timeStep)-3));
            upFactor = a*(b+c);
            downFactor = a*(b-c);
            upProbability = (exp(riskFreeRate*timeStep) - downFactor) / (upFactor - downFactor);
        }
    };

    class Trigeorgis
    {
    public:
        static void factors(double& upFactor,
                            double& downFactor,
                            double& upProbability,
                            double const volatility,
                            double const riskFreeRate,
                            double const timeStep)
        {
            double const volSquared(volatility*volatility);
            double const y(riskFreeRate - 0.5*volSquared);
            double const x(sqrt(timeStep*(volSquared + y*y*timeStep)));
            upFactor = exp(x);
            downFactor = 1.0 / upFactor;
            upProbability = 0.5*(1+y*timeStep/x);
        }
    };

    template <typename FACTOR_MODEL=CoxRossRubinstein>
    class BinomialTree
    {
    public:
        class Evaluation
        {
        public:
            Evaluation(std::size_t const initHeight=25);
            ~Evaluation();

            double delta() const { return _delta; }
            double gamma() const { return _gamma; }
            double theta() const { return _theta; }
            
            double theoreticalPrice() const { return _theo; }
        private:
            friend class BinomialTree<FACTOR_MODEL>;
        
            void validateHeight(std::size_t const height)
            {
                if( height > _maxHeight ) {
                    if( _underlyerPrices ) delete [] _underlyerPrices;
                    if( _values ) delete _values;
                    _underlyerPrices = new double[height+1];
                    _values = new double[height+1];
                    _maxHeight = height;
                }
            }

            std::size_t _maxHeight;
            double *_underlyerPrices;
            double *_values;
            double _theo;
            double _delta;
            double _gamma;
            double _theta;
        };

        BinomialTree() { }
        ~BinomialTree() { }

        bool valuate(Evaluation& eval,
                     VanillaOption const& option,
                     double const underlyerPrice,
                     double const volatility,
                     double const riskFreeRate,
                     int const height)
        {
            eval.validateHeight(height);
            if( option.type() == VanillaOption::CALL ) {
                if( option.style() == VanillaOption::AMERICAN ) {
                    return americanCall(eval, option, underlyerPrice,
                                        volatility, riskFreeRate, height);
                } else {
                    return europeanCall(eval, option, underlyerPrice,
                                        volatility, riskFreeRate, height);
                }
            } else {
                if( option.style() == VanillaOption::AMERICAN ) {
                    return americanPut(eval, option, underlyerPrice,
                                       volatility, riskFreeRate, height);
                } else {
                    return europeanPut(eval, option, underlyerPrice,
                                       volatility, riskFreeRate, height);
                }
            }
        }
    private:
        bool americanCall(Evaluation& eval,
                          VanillaOption const& option,
                          double const underlyerPrice,
                          double const volatility,
                          double const riskFreeRate,
                          int const height);
        bool europeanCall(Evaluation& eval,
                          VanillaOption const& option,
                          double const underlyerPrice,
                          double const volatility,
                          double const riskFreeRate,
                          int const height);
        bool americanPut(Evaluation& eval,
                         VanillaOption const& option,
                         double const underlyerPrice,
                         double const volatility,
                         double const riskFreeRate,
                         int const height);
        bool europeanPut(Evaluation& eval,
                         VanillaOption const& option,
                         double const underlyerPrice,
                         double const volatility,
                         double const riskFreeRate,
                         int const height);
    };
}

/***
 * Implementation
 ***/
template <typename FACTOR_MODEL>
provident::BinomialTree<FACTOR_MODEL>::Evaluation::Evaluation(std::size_t const initHeight)
:   _maxHeight(0),
    _underlyerPrices(0),
    _values(0)
{
    validateHeight(initHeight);
}

template <typename FACTOR_MODEL>
provident::BinomialTree<FACTOR_MODEL>::Evaluation::~Evaluation()
{
    if( _underlyerPrices ) delete [] _underlyerPrices;
    if( _values ) delete [] _values;
}

template <typename FACTOR_MODEL>
bool
provident::BinomialTree<FACTOR_MODEL>::americanCall(provident::BinomialTree<FACTOR_MODEL>::Evaluation& eval,
                                         provident::VanillaOption const& option,
                                         double const underlyerPrice,
                                         double const volatility,
                                         double const riskFreeRate,
                                         int const height)
{
    double const timeDelta(option.expiryYears() / height);
    double const stepRate(exp(riskFreeRate*timeDelta));
    double const stepRateInv(1.0 / stepRate);

    double u, d, pUp;
    FACTOR_MODEL::factors(u,d,pUp,volatility,riskFreeRate,timeDelta);
    double const dd(d*d);
    double const pDown(1.0 - pUp);

    eval._underlyerPrices[0] = underlyerPrice * pow(u,height);
    eval._values[0] = std::max(0.0,eval._underlyerPrices[0]-option.strikePrice());
    //printf("up[%d]=%f, strike=%f, val[%d]=%f\n", 0, eval._underlyerPrices[0],
    //       option.strikePrice(), 0, eval._values[0]);
    int ii = 0;
    while(ii <= height and eval._values[ii++] != 0.0) {
        eval._underlyerPrices[ii] = dd * eval._underlyerPrices[ii-1];
        eval._values[ii] = std::max(0.0,eval._underlyerPrices[ii]-option.strikePrice());
        //printf("up[%d]=%f, strike=%f, val[%d]=%f\n", ii, eval._underlyerPrices[ii],
        //       option.strikePrice(), ii, eval._values[ii]);
    }
    while( ii <= height ) {
        eval._values[ii++] = 0.0;
    }

    for(int step = height; step; --step) {
        for(ii = 0; ii != step; ++ii) {
            if( eval._values[ii] ) {
                eval._values[ii] = (pUp*eval._values[ii]+pDown*eval._values[ii+1]) * stepRateInv;
            } else {
                 break;
            }
            //printf("eval._values[%d] = %f\n", ii, eval._values[ii]);
            //eval._underlyerPrices[ii] = d*eval._underlyerPrices[ii+1];
            //eval._values[ii] = std::max(eval._values[ii],eval._underlyerPrices[ii]-option.strikePrice());
        }
    }
    eval._theo = eval._values[0];
    return true;
}

template <typename FACTOR_MODEL>
bool
provident::BinomialTree<FACTOR_MODEL>::europeanCall(provident::BinomialTree<FACTOR_MODEL>::Evaluation& eval,
                                                    provident::VanillaOption const& option,
                                                    double const underlyerPrice,
                                                    double const volatility,
                                                    double const riskFreeRate,
                                                    int const height)
{
    double const timeDelta(option.expiryYears() / height);
    double const stepRate(exp(riskFreeRate*timeDelta));
    double const stepRateInv(1.0 / stepRate);
    
    double u, d, pUp;
    FACTOR_MODEL::factors(u,d,pUp,volatility,riskFreeRate,timeDelta);
    double const dd(d*d);
    double const pDown(1.0 - pUp);

    eval._underlyerPrices[0] = underlyerPrice * pow(u,height);
    eval._values[0] = std::max(0.0,eval._underlyerPrices[0]-option.strikePrice());
    //printf("up[%d]=%f, strike=%f, val[%d]=%f\n", 0, eval._underlyerPrices[0],
    //       option.strikePrice(), 0, eval._values[0]);
    int ii = 0;
    while(ii <= height and eval._values[ii++] != 0.0) {
        eval._underlyerPrices[ii] = dd * eval._underlyerPrices[ii-1];
        eval._values[ii] = std::max(0.0,eval._underlyerPrices[ii]-option.strikePrice());
        //printf("up[%d]=%f, strike=%f, val[%d]=%f\n", ii, eval._underlyerPrices[ii],
        //       option.strikePrice(), ii, eval._values[ii]);
    }
    while( ii <= height ) {
        eval._values[ii++] = 0.0;
    }

    for(int step = height; step; --step) {
        for(ii = 0; ii != step; ++ii) {
            if( eval._values[ii] ) {
                eval._values[ii] = (pUp*eval._values[ii]+pDown*eval._values[ii+1]) * stepRateInv;
            } else {
                 break;
            }
            //printf("eval._values[%d] = %f\n", ii, eval._values[ii]);
            //eval._underlyerPrices[ii] = d*eval._underlyerPrices[ii+1];
            //eval._values[ii] = std::max(eval._values[ii],eval._underlyerPrices[ii]-option.strikePrice());
        }
    }
    eval._theo = eval._values[0];
    return true;
}

template <typename FACTOR_MODEL>
bool
provident::BinomialTree<FACTOR_MODEL>::americanPut(provident::BinomialTree<FACTOR_MODEL>::Evaluation& eval,
                                                   provident::VanillaOption const& option,
                                                   double const underlyerPrice,
                                                   double const volatility,
                                                   double const riskFreeRate,
                                                   int const height)
{
    double const timeDelta(option.expiryYears() / height);
    double const stepRate(exp(riskFreeRate*timeDelta));
    double const stepRateInv(1.0 / stepRate);
    
    double u, d, pUp;
    FACTOR_MODEL::factors(u,d,pUp,volatility,riskFreeRate,timeDelta);
    double const uu(u*u);
    double const pDown(1.0 - pUp);

    //printf("u=%f, d = %f\n", u, d);
    eval._underlyerPrices[0] = underlyerPrice * pow(d,height);
    eval._values[0] = std::max(0.0,option.strikePrice()-eval._underlyerPrices[0]);
    //printf("up[%d]=%f, strike=%f, val[%d]=%f\n", 0, eval._underlyerPrices[0],
    //       option.strikePrice(), 0, eval._values[0]);
    int ii = 0;
    while(ii <= height and eval._values[ii++] != 0.0) {
        eval._underlyerPrices[ii] = uu * eval._underlyerPrices[ii-1];
        eval._values[ii] = std::max(0.0,option.strikePrice()-eval._underlyerPrices[ii]);
        //printf("up[%d]=%f, strike=%f, val[%d]=%f\n", ii, eval._underlyerPrices[ii],
        //       option.strikePrice(), ii, eval._values[ii]);
    }
    while( ii <= height ) {
        eval._values[ii++] = 0.0;
    }

    for(int step = height; step; --step) {
        for(ii = 0; ii != step; ++ii) {
            if( eval._values[ii] ) {
                eval._values[ii] = (pDown*eval._values[ii]+pUp*eval._values[ii+1]) * stepRateInv;
                eval._underlyerPrices[ii] *= u;//= d*eval._underlyerPrices[ii+1];
                //printf("eval._values[%d] = %f, vs. value in exercising: %f\n", ii, eval._values[ii],
                //       option.strikePrice() - eval._underlyerPrices[ii]);
                eval._values[ii] = std::max(eval._values[ii],option.strikePrice()-eval._underlyerPrices[ii]);
            } else {
                 break;
            }
        }
    }
    eval._theo = eval._values[0];
    return true;
}

template <typename FACTOR_MODEL>
bool
provident::BinomialTree<FACTOR_MODEL>::europeanPut(provident::BinomialTree<FACTOR_MODEL>::Evaluation& eval,
                                                   provident::VanillaOption const& option,
                                                   double const underlyerPrice,
                                                   double const volatility,
                                                   double const riskFreeRate,
                                                   int const height)
{
    double const timeDelta(option.expiryYears() / height);
    double const stepRate(exp(riskFreeRate*timeDelta));
    double const stepRateInv(1.0 / stepRate);
    
    double u, d, pUp;
    FACTOR_MODEL::factors(u,d,pUp,volatility,riskFreeRate,timeDelta);
    double const uu(u*u);
    double const pDown(1.0 - pUp);

    eval._underlyerPrices[0] = underlyerPrice * pow(d,height);
    eval._values[0] = std::max(0.0,option.strikePrice()-eval._underlyerPrices[0]);
    //printf("up[%d]=%f, strike=%f, val[%d]=%f\n", 0, eval._underlyerPrices[0],
    //       option.strikePrice(), 0, eval._values[0]);
    int ii = 0;
    while(ii <= height and eval._values[ii++] != 0.0) {
        eval._underlyerPrices[ii] = uu * eval._underlyerPrices[ii-1];
        eval._values[ii] = std::max(0.0,option.strikePrice()-eval._underlyerPrices[ii]);
        //printf("up[%d]=%f, strike=%f, val[%d]=%f\n", ii, eval._underlyerPrices[ii],
        //       option.strikePrice(), ii, eval._values[ii]);
    }
    while( ii <= height ) {
        eval._values[ii++] = 0.0;
    }

    for(int step = height; step; --step) {
        for(ii = 0; ii != step; ++ii) {
            if( eval._values[ii] ) {
                eval._values[ii] = (pDown*eval._values[ii]+pUp*eval._values[ii+1]) * stepRateInv;
                //eval._underlyerPrices[ii] = u*eval._underlyerPrices[ii+1];
                //eval._values[ii] = std::max(eval._values[ii],option.strikePrice()-eval._underlyerPrices[ii]);
                //printf("eval._values[%d] = %f\n", ii, eval._values[ii]);
            } else {
                 break;
            }
        }
    }
    eval._theo = eval._values[0];
    return true;
}

/*template <typename FACTOR_MODEL>
bool
BinomialTree<FACTOR_MODEL>::europeanCall(Evaluation& eval,
                                         VanillaOption const& option,
                                         double const underlyerPrice,
                                         double const volatility,
                                         double const riskFreeRate,
                                         int const height)
{
    //double values[200];
    double* const values(eval._values);

    double const timeDelta(option.expiryYears() / height);
    double const stepRate(exp(riskFreeRate*timeDelta));
    double const stepRateInv(1.0 / stepRate);
    double const u(exp(volatility * sqrt(timeDelta)));
    double const d(1.0 / u);
    double const pUp((stepRate - d) / (u - d));
    double const pDown(1.0 - pUp);

    //printf("stepRate=%g, u=%g, d=%g, pUp=%g, pDown=%g\n",
    //        stepRate, u, d, pUp, pDown);

    //for(int level = 0; level <= height; ++level) {
    //    for(int ups = 0; ups <= level; ++ups) {
    //        printf("S[%d,%d]=%g\n", ups, level-ups,
    //               underlyerPrice * pow(u,ups) * pow(d,level-ups));
    //    }
   // }
    
    // Terminal values
    double uFactor(1);
    double dFactor(pow(d,height));
    for(int ii = 0; ii <= height; ++ii) {
        double const Sprice(underlyerPrice * uFactor * dFactor);
        //printf("%d ups spot price=%g\n", ii, Sprice);
        values[ii] = std::max(0.0,Sprice - option.strikePrice());
        //printf("Init value[%i]=%g\n", ii, values[ii]);
        uFactor *= u;
        dFactor *= u; 
    }

    // Work backwards
    double valueU(0), valueD(0), valueUU(0), valueUD(0), valueDD(0);
    for(int jj = height; jj >= 0; --jj) {
        for(int ii = 0; ii != jj; ++ii) {
            if( values[ii] != 0 or values[ii+1] != 0 ) {
                values[ii] = (pUp*values[ii+1] + pDown*values[ii]) * stepRateInv;
            } else {
                values[ii] = 0;
            }
            //printf("Set values[%d]=%g\n", ii, values[ii]);
        }
        if( jj == 2 ) {
            valueU = values[1];
            valueD = values[0];
            //printf("set 1\n");
        } else if( jj == 3 ) {
            valueUU = values[2];
            valueUD = values[1];
            valueDD = values[0];
            //printf("set 2\n");
        }
    }

    eval._theo = values[0];

    // Delta
    double const spotU(underlyerPrice * u);
    double const spotD(underlyerPrice * d);
    eval._delta = (valueU - valueD) / (spotU - spotD);
    //printf("delta = (%g - %g) / (%g - %g) = %g\n",
    //       valueU, valueD, spotU, spotD, eval._delta);

    // Gamma
    double const spotUU(spotU * u);
    double const spotUD(spotU * d);
    double const spotDD(spotD * d);
    double const deltaU((valueUU - valueUD) / (spotUU - spotUD));
    double const deltaD((valueUD - valueDD) / (spotUD - spotDD));
    eval._gamma = (deltaU - deltaD) / (spotU - spotD);
    
    // Theta
    //eval._theta = (valueUD - eval._theo) / (2 * timeDelta * 365);
    eval._theta = (riskFreeRate*eval._theo) - (riskFreeRate*underlyerPrice*eval._delta + 0.5*pow(volatility,2)*pow(underlyerPrice,2)*eval._gamma);
    eval._theta /= 365;

    return true;
}

template <typename FACTOR_MODEL>
bool
BinomialTree<FACTOR_MODEL>::americanPut(Evaluation& eval,
                                        VanillaOption const& option,
                                        double const underlyerPrice,
                                        double const volatility,
                                        double const riskFreeRate,
                                        int const height)
{
    return false;
}

template <typename FACTOR_MODEL>
bool
BinomialTree<FACTOR_MODEL>::europeanPut(Evaluation& eval,
                                        VanillaOption const& option,
                                        double const underlyerPrice,
                                        double const volatility,
                                        double const riskFreeRate,
                                        int const height)
{
    return false;
}*/


#endif /* PROVIDENT_CRR_BINOMIAL_HPP_ */

